#include<vector>
#include<iostream>
#include<sstream>
#include<fstream>
#include<openssl/aes.h>
#include<openssl/rand.h>
#include<stdint.h>
#include<string.h>
#include<time.h>
#include<sys/time.h>
#include<stdlib.h>

#include "utils.h"
#include "streamforce_crypto_KPAbe.h"

using namespace std;


const unsigned int AES_KEY_SIZE=128;
const unsigned int AES_BLOCK_BITS=128;
const unsigned int AES_BLOCK_BYTES=AES_BLOCK_BITS/8;
const unsigned int AES_KEY_BYTES=AES_KEY_SIZE/8;

const string defaultKey="Lorem ipsum dolor sit amet, consectetur adipiscing elit. And hello world here";	
const string defaultOPEKey = "Default OPE";


string
randomBytes(unsigned int len)
{
    string s;
    s.resize(len);
    RAND_bytes((uint8_t*) &s[0], len);
    return s;
}

JNIEXPORT jstring JNICALL Java_streamforce_crypto_KPAbe_bytesToString
  (JNIEnv *env, jobject obj, jbyteArray bytes){
	int len = env->GetArrayLength(bytes);
	unsigned char* bArray = (unsigned char*)malloc(len*sizeof(char));
	env->GetByteArrayRegion(bytes,0,len,(jbyte*)bArray);
	string res = marshallByteArray(bArray,len);
	stringstream ss;
	ss << res;
	return env->NewStringUTF(&(ss.str())[0]);
}


string
marshallBinary(const string &s)
{
    string r;
        r += "X\'";

    const char *sp = &s[0];
    static const char *digits = "0123456789ABCDEF";
    size_t l = s.length();
    for (size_t i = 0; i < l; i++) {
        r += digits[sp[i] >> 4 & 0xf];
        r += digits[sp[i] & 0xf];
    }

        r += "\'";
    return r;
}

static unsigned char
getFromHex(const string &hexValues)
{
    unsigned int v;
    sscanf(hexValues.c_str(), "%2x", &v);
    return (unsigned char) v;
}

string marshallByteArray(byte* data, int length){
	 string r;
  r+= "X\'";

  static const char *digits = "0123456789ABCDEF";
  for (int i=0; i<length; i++){
    r+= digits[data[i] >> 4 & 0xf];
    r+= digits[data[i] & 0xf];
  }
  r+= "\'";
  return r;
}

string
unmarshallBinary(const string &s)
{
    int offset;
    size_t len = s.length();

    offset = 2;
    len = len - 1;     // removing last apostrophe

    string r;
    for (int i = 0; i < (len-offset)/2; i++)
        r += getFromHex(&s[offset+i*2]);
    return r;
}

AES_KEY *
get_AES_enc_key(const string &key)
{
    char *tempkey = "Lorem ipsum dolor sit amet, consectetur adipiscing elit.";	
    AES_KEY * aes_key = new AES_KEY();
    //AES_set_encrypt_key((const uint8_t*) key.c_str(), AES_KEY_SIZE,
	//		aes_key);
	AES_set_encrypt_key((unsigned char*)tempkey, 128,aes_key);
    return aes_key;
}

AES_KEY *
get_default_AES_enc_key()
{
    AES_KEY * aes_key = new AES_KEY();
    AES_set_encrypt_key((unsigned char*)defaultKey.data(), AES_KEY_SIZE, aes_key);
	return aes_key;
}

AES_KEY *
get_default_AES_dec_key()
{
    AES_KEY * aes_key = new AES_KEY();
    AES_set_decrypt_key((const uint8_t*) defaultKey.c_str(), AES_KEY_SIZE,
			aes_key);
    return aes_key;
}

string get_default_OPE_key()
{
  return defaultOPEKey;
}

AES_KEY *
get_AES_dec_key(const string &key)
{

    AES_KEY * aes_key = new AES_KEY();
    AES_set_decrypt_key((const uint8_t*) key.c_str(), AES_KEY_SIZE,
                            aes_key);
    return aes_key;

}

static vector<unsigned char>
getIVec(string salt)
{
    vector<unsigned char> ivec(AES_BLOCK_BYTES);

    memcpy(&ivec[0], salt.data(), min(salt.length(), (size_t) AES_BLOCK_BYTES));

    return ivec;
}

template<typename SIZE_T>
SIZE_T
getBlocks(unsigned int unit, SIZE_T len) {
    SIZE_T blocks = len / unit;
    if (len > blocks * unit) {
        blocks++;
    }
    return blocks;
}

vector<unsigned char>
pad(vector<unsigned char> data, unsigned int unit)
{
    size_t blocks = getBlocks(unit, data.size());
    size_t multipleLen = blocks * unit;
    size_t padding;
    if (multipleLen == data.size()) {
        padding = unit;
    } else {
        padding = multipleLen - data.size();
    }
    size_t paddedLen = data.size() + padding;

    vector<unsigned char> res(paddedLen, 0);
    res[paddedLen-1] = (unsigned char)padding;
    memcpy(&res[0], &data[0], data.size());
    return res;
}

vector<unsigned char>
unpad(vector<unsigned char> data)
{
    size_t len = data.size();
    size_t actualLen = len - (int)data[len-1];
    vector<unsigned char> res(actualLen);
    memcpy(&res[0], &data[0], actualLen);
    return res;
}

void 
encrypt_AES_CBC(const string &ptext, unsigned char **ctext, int *len, const AES_KEY * enckey, string salt, bool dopad)
{
    vector<unsigned char> ptext_buf;
    if (dopad) {
        ptext_buf = pad(vector<unsigned char>(ptext.begin(), ptext.end()), AES_BLOCK_BYTES);
    } else {
        ptext_buf = vector<unsigned char>(ptext.begin(), ptext.end());
    }
    //vector<unsigned char> ctext_buf = vector<unsigned char>(ptext_buf.size());
    *len = ptext_buf.size();
    *ctext = (unsigned char*) malloc((*len)*sizeof(unsigned char));
	
    vector<unsigned char> ivec = getIVec(salt);
    //AES_cbc_encrypt(&ptext_buf[0], &ctext_buf[0], ptext_buf.size(), enckey, &ivec[0], AES_ENCRYPT);
    AES_cbc_encrypt(&ptext_buf[0], *ctext, ptext_buf.size(), enckey, &ivec[0], AES_ENCRYPT);

    //return string((char *) &ctext_buf[0], ctext_buf.size());
}

string
decrypt_AES_CBC(const string &ctext, const AES_KEY * deckey, string salt, bool dounpad)
{
    vector<unsigned char> ptext_buf(ctext.size());
    vector<unsigned char> ivec = getIVec(salt);

    AES_cbc_encrypt((const unsigned char *) ctext.data(), &ptext_buf[0], ctext.size(), deckey, &ivec[0], AES_DECRYPT);


    if (dounpad) {
      vector<unsigned char> res = unpad(ptext_buf);
       
        return string((char *)&res[0], res.size());
    } else {
        return string((char *)&ptext_buf[0], ptext_buf.size());
    }
}

string
stringToByteInts(const string &s)
{
    stringstream ss;
    bool first = true;
    for (size_t i = 0; i < s.length(); i++) {
        if (!first)
            ss << " ";
        ss << (unsigned int) (uint8_t) s[i];
        first = false;
    }
    return ss.str().data();
}

int user_initDiscreteLogMap(int max, int streamId, kpabe_pub_t* pub){
	struct timeval start, end;
	gettimeofday(&start,NULL); 	

	element_t exp;
	element_init_Zr(exp,pub->p);
	element_t val;
	element_init_GT(val,pub->p);
	dLogMap lm;
	ptMap pm; 
	for (int i=0; i<max; i++){
		element_set_si(exp,i);
		element_pow_zn(val,pub->gt,exp);
		element_wrap val1;
		element_init_GT(val1.element,pub->p);
		element_set(val1.element,val);
		(lm.map)[val1] = i;
		(pm.map)[i] = val1; 	
	}
	user_discreteLogMap[streamId]=lm;
	user_plaintextMap[streamId]=pm;
	gettimeofday(&end,NULL);
	cout << "user discrete log init time = " << (end.tv_sec-start.tv_sec)*1000 + (end.tv_usec - start.tv_usec)/1000 << endl;

 	return 0;	
}

int owner_initDiscreteLogMap(int max, int streamId, kpabe_pub_t* pub){
	struct timeval start, end;
	gettimeofday(&start,NULL); 	
	element_t exp;
	element_init_Zr(exp,pub->p);
	element_t val;
	element_init_GT(val,pub->p);
	dLogMap lm;
	ptMap pm; 
	for (int i=0; i<max; i++){
		element_set_si(exp,i);
		element_pow_zn(val,pub->gt,exp);
		element_wrap val1;
		element_init_GT(val1.element,pub->p);
		element_set(val1.element,val);
		(lm.map)[val1] = i;
		(pm.map)[i] = val1; 	
	}
	owner_discreteLogMap[streamId]=lm;
	owner_plaintextMap[streamId]=pm;
	gettimeofday(&end,NULL);
	cout << "owner discrete log init time = " << (end.tv_sec-start.tv_sec)*1000 + (end.tv_usec - start.tv_usec)/1000 << endl;
 	return 0;	
}

gint str_cmp(gconstpointer a, gconstpointer b){
	return strcmp((const char*)a,(const char*)b);
}

char** parseAttributeUniverse(char** argv, int size){
        GSList* alist;
        GSList* ap;
        alist = 0;
        ap = 0;
        for (int i=0; i<size; i++)
                parse_attribute_universe(&alist,argv[i]);
        alist = g_slist_sort(alist,str_cmp);
        int n = g_slist_length(alist);
        char** atts = (char**)malloc((n+1)*sizeof(char*));
        int i=0;
        for (ap = alist; ap; ap = ap->next)
                atts[i++] = (char*)ap->data;
        atts[i] = 0;
	return atts;
}

char** parseEncryptionAttributes(char** argv, int size){
	GSList* alist;
	GSList* ap;
	alist = 0;
	ap = 0;
	for (int i=0; i<size; i++)
		parse_attribute(&alist,argv[i]);
	alist = g_slist_sort(alist,str_cmp);
	int n = g_slist_length(alist);
	char** encAtts = (char**)malloc((n+1)*sizeof(char*));
	int i=0;
	for (ap = alist; ap; ap = ap->next)
		encAtts[i++] = (char*)ap->data;
	encAtts[i] = 0;
	return encAtts;
}

int toByteArray(GByteArray* buf, unsigned char** byteArray, int *len){
	*len = buf->len;
	*byteArray = (unsigned char*)malloc((*len)*sizeof(char));
	for (int i=0; i<(*len); i++)
		(*byteArray)[i] = buf->data[i];
	return 0;
}

kpabe_pub_t* convertToPubKey(unsigned char* pubBytes, int pubLen){
	GByteArray* buf = convertToByteArray(pubBytes, pubLen);
	return kpabe_pub_unserialize(buf, 1);
}

kpabe_prv_t* convertToPrvKey(kpabe_pub_t* pub, unsigned char* prvBytes, int prvLen){
	return kpabe_prv_unserialize(pub, convertToByteArray(prvBytes, prvLen), 1);
}

//the element is in Zn
int convertToElement(kpabe_pub_t* pub, unsigned char *bytes, int len, element_t e){
	element_init_Zr(e,pub->p);
	len = element_from_bytes(e,bytes);
	return 0;	
}

//this is used by the user 
int toPlaintext_simple(int streamId, element_t t){
	element_wrap et;
	kpabe_pub_t* pub = userPKMap[streamId]; 

	element_init_GT(et.element,pub->p);

	element_set(et.element,t);

	unordered_map<element_wrap, int, elementHash, equalElement>::const_iterator it =
	user_discreteLogMap[streamId].map.find(et);
	
	if (it==user_discreteLogMap[streamId].map.end()){
		return -1;
	}
	return user_discreteLogMap[streamId].map[et];

}

int toPlaintext_slidingWindow(int streamId, element_t t, element_t mask){
	element_wrap et;
	kpabe_pub_t* pub = userPKMap[streamId];
	element_init_GT(et.element,pub->p);

	element_t unmask, pt;
	element_init_GT(unmask, pub->p);
	element_init_GT(pt,pub->p);
	element_pow_zn(unmask,pub->gt,mask);
	//element_div(pt,t,unmask);
	element_mul(pt,t,unmask);	

	element_set(et.element,pt);

	unordered_map<element_wrap, int, elementHash, equalElement>::const_iterator it =
	user_discreteLogMap[streamId].map.find(et);

	if (it==user_discreteLogMap[streamId].map.end()){
		return -1;
	}
	return user_discreteLogMap[streamId].map[et];

}

GByteArray* convertToByteArray(unsigned char* cipher, int len){
	GByteArray* buf = g_byte_array_new();
	g_byte_array_set_size(buf, len);
	for (int i=0; i<len; i++){
		buf->data[i] = cipher[i];
	}		
	return buf;
}

int compare_element_zn(kpabe_pub_t* pub, unsigned char *e1, unsigned char *e2){
	element_t x1, x2;
	element_init_Zr(x1, pub->p);
	element_init_Zr(x2, pub->p);
	element_from_bytes(x1,e1);
	element_from_bytes(x2,e2);
	return element_cmp(x1,x2);
}
